from ultralytics import YOLO
import cv2
import time
import numpy as np
from image_provider import ImageProvider

SEG_MODEL_PATH = "./merge_yolo11m_seg_rknn_model"
POSE_MODEL_PATH = "./merge_yolo11m_pose_rknn_model"

# 加载模型
seg_model = YOLO(SEG_MODEL_PATH, task="segment")
pose_model = YOLO(POSE_MODEL_PATH, task="pose")

# 加载测试图像
img1 = cv2.imread("./data/Color1/330.png")
img2 = cv2.imread("./data/Color1/340.png")
img3 = cv2.imread("./data/Color1/350.png")
images = [img1, img2, img3]

# 添加计时
start_time = time.time()  # 记录开始时间

# 执行推理
result = seg_model.predict(img1, verbose=False,batch=3)  # 确保verbose=False，关闭调试信息输出。1)  # 流推理

end_time = time.time()  # 记录结束时间
elapsed_time = end_time - start_time  # 计算耗时

print(f"推理耗时: {elapsed_time:.4f} 秒")